import os
import time
import numpy as np
import pickle
import struct
import re
import cv2
from tqdm import tqdm
import imageio
import json
from PIL import Image
from multiprocessing import Pool
# import threadpool
# import threading
from collections import defaultdict
import random
import copy
from ctypes import CDLL,byref,create_string_buffer,cdll
from PIL import Image, ImageColor, ImageFont, ImageDraw, ImageFilter
import sys
import imgaug as ia
import imgaug.augmenters as iaa
import imgaug.parameters as iap
from imgaug.augmentables.polys import Polygon, PolygonsOnImage


from tool import filesystem,imgaug_tool,utils, data_enhance, via_tool # export PYTHONPATH=$PYTHONPATH:`pwd`



PLATE_CHARS_PROVINCE = ["京", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑",
                        "苏", "浙", "皖", "闽", "赣", "鲁", "豫", "鄂", "湘", "粤",
                        "桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁",
                        "新"]
# PLATE_CHARS_DIGIT = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
# PLATE_CHARS_LETTER = ["A", "B", "C", "D", "E", "F", "G",
#                         "H", "J", "K", "L", "M", "N",
#                         "P", "Q", "R", "S", "T",
#                         "U", "V", "W", "X", "Y", "Z"]

PLATE_CHARS_DIGIT = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"] \
                    + ["0"] * 3 \
                    + ["2"] * 3 \
                    + ["6"] * 3 \
                    + ["8"] * 3 \
                    + ["1"] * 3  

PLATE_CHARS_LETTER = ["A", "B", "C", "D", "E", "F", "G",
                        "H", "J", "K", "L", "M", "N",
                        "P", "Q", "R", "S", "T",
                        "U", "V", "W", "X", "Y", "Z"] \
                    + ["Q"] * 3 \
                    + ["Z"] * 3 \
                    + ["G"] * 3 \
                    + ["B"] * 3 \
                    + ["J"] * 2 \
                    + ["U"] * 2 \
                    + ["T"] * 2      


provinces = ["皖", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑", "苏", "浙", "京", "闽", "赣", "鲁", "豫", "鄂", "湘", "粤", "桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁", "新", "警", "学", "O"]
alphabets = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'O']
ads = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'O']


# def init_sort_cpp():
#     # so_path = os.path.dirname(os.path.abspath(__file__)) + os.sep + "libsort_rect.so"     
#     so_path = "/home/swls/work_dir/git/rect_sort_cpp/build/src/libsort_rect.so"
#     rect_sort_so = cdll.LoadLibrary(so_path)
#     return rect_sort_so

# # rect_sort_so = init_sort_cpp()



# def get_point_cpp(points):

#     info_dict = []
#     for point in points:
#         i_d = {
#             "x":int(point[0]),
#             "y":int(point[1]),
#         }
#         info_dict.append(i_d)
#     input_info = json.dumps(info_dict)
#     # print(input_info)

#     so_path_buffer = create_string_buffer(input_info.encode('utf-8'), 65536*100)    
#     output_msg = create_string_buffer(65536*100)

#     out = rect_sort_so.get_point(byref(so_path_buffer),byref(output_msg))  
    
#         # print(output_msg.value.decode('utf-8'))
#     out_put_info = json.loads(output_msg.value.decode('utf-8'))

#     # return 0

#     new_points = [[out_put_info['x1'], out_put_info['y1']],
#                 [out_put_info['x2'], out_put_info['y2']],
#                 [out_put_info['x3'], out_put_info['y3']],
#                 [out_put_info['x4'], out_put_info['y4'] ]]

#     return new_points



def iou(rect1, rect2):
    y0 = np.max([rect1[1],rect2[1]])
    y1 = np.min([rect1[1] + rect1[3],rect2[1]+rect2[3]])

    x0 = np.max([rect1[0],rect2[0]])
    x1 = np.min([rect1[0] + rect1[2],rect2[0]+rect2[2]])
    return x1 > x0 and y1 > y0


def GenNewCh(f,val, w=45, h=90, color = (0,0,0)):
    img=Image.new("RGBA", (w,h),(0,0,0,0))
    draw = ImageDraw.Draw(img)
    draw.text((0, 0),val,color,font=f)
    return img

class TextImageGenerator:
    def __init__(self, current_name, current_type):
        self.plate_type = current_type
        self.current_name = current_name

        if self.plate_type in [ 0, 2, 3, 4, 5]:
            # # 128 测定值
            self.fontCE = ImageFont.truetype(fontChEng[self.plate_type], 128, 0)    # 省简称使用字体
            # self.pil_img = Image.new("RGBA", (440, 140), (0, 0, 0,0))

        elif self.plate_type == 1:
            # # 124 测定值
            self.fontCE = ImageFont.truetype(fontChEng[self.plate_type], 124, 0)    # 省简称使用字体
            # self.pil_img = Image.new("RGBA", (480, 140), (0, 0, 0,0))
        else:
            raise NotImplementedError
        
        self.template_image_paths = [TEMPLATE_IMAGES[current_type] + os.sep + f for f in os.listdir(TEMPLATE_IMAGES[current_type])]

    # 生成文本行图片
    def gen_text_line_image(self, pil_img, text, w=45, h=90):
        """
            # 蓝牌字体本身有上下间距
            # 绿牌字体没有上下左右间距

            # 整行一个框 用于训练DB
        """
        pts = []
        labels = []
        offset = 0
        # string_length = 8 if self.plate_type == 1 else 7
        string_length = 7
        if self.plate_type == 1:
            color = (5,5,5)
            first_points = [0,0,0,0]

            pil_img.paste(GenNewCh(self.fontCE, text[0], 45, 90, color), (offset+15, 25))
            first_points[0] = [offset+15, 25]
            first_points[3] = [offset+15, 25+90]

            offset += 15+45
            for i in range(string_length-1):
                if i == 1:
                    offset += 40 
                offset += 9 
                if text[i+1] != "x":
                    w = 43 
                    pil_img.paste( GenNewCh(self.fontCE, text[i+1], w, 90, color), (offset, 25))

                offset += 43 

                if i == string_length -2:
                    first_points[1] = [offset, 25]
                    first_points[2] = [offset, 25+90]

            pts.append(first_points)
            labels.append(text)

        elif self.plate_type in [ 0, 2 ]:
            first_points = [0,0,0,0]
            w = 45

            color = (252,252,252) if self.plate_type == 0 else (5,5,5)
            pil_img.paste(GenNewCh(self.fontCE, text[0], 45, 90, color), (offset+15, 25))
            first_points[0] = [offset+15, 25]
            first_points[3] = [offset+15, 25+90]

            offset += 15+45
            for i in range(string_length-1):
                if i == 1:
                    offset +=  22
                offset +=  12
                if text[i+1] != "x":
                    pil_img.paste( GenNewCh(self.fontCE, text[i+1], w, 90, color), (offset, 25))

                offset += 45

                if i == string_length -2:
                    first_points[1] = [offset, 25]
                    first_points[2] = [offset, 25+90]

            pts.append(first_points)
            labels.append(text)

        elif self.plate_type in [ 3 ]:
            color = (5,5,5)

            first_points = [0, 0, 0, 0]
            first_text = ""
            offset = 110
            for i in range(2):
                img_1 = GenNewCh(self.fontCE, text[i], 45, 90, color)
                img_1 = img_1.resize((80, 60), Image.ANTIALIAS)
                pil_img.paste(img_1, (offset, 15))
                if i == 0:
                    first_points[0] = [offset, 15]
                    first_points[3] = [offset, 15+60]
                else :
                    first_points[1] = [offset+80, 15]
                    first_points[2] = [offset+80, 15+60]
                first_text += text[i]
                offset += 140

            pts.append(first_points)
            labels.append(first_text)

            second_points = [0, 0, 0, 0]
            second_text = ""
            offset = 27
            for i in range(2,string_length):
                w = 65
                img = GenNewCh(self.fontCE, text[i], 45, 90, color)
                img = img.resize((w, 110), Image.ANTIALIAS)
                pil_img.paste( img, (offset, 90))

                if i == 2:
                    second_points[0] = [offset, 90]
                    second_points[3] = [offset, 90+110]
                elif i == string_length - 1:
                    second_points[1] = [offset + w, 90]
                    second_points[2] = [offset + w, 90+110]
                second_text += text[i]

                offset += 15 + 65
            
            pts.append(second_points)
            labels.append(second_text)

        elif self.plate_type in [ 4, 5 ]:
            # 4 -> 黑牌
            color = (252,252,252) if self.plate_type == 4 else (5,5,5)
            w = 45

            first_points = [0,0,0,0]

            if text[0] == "使":         # 红色
                pil_img.paste(GenNewCh(self.fontCE, text[0], w, 90, (255, 0 ,0)), (offset+15, 25))
            else:
                pil_img.paste(GenNewCh(self.fontCE, text[0], w, 90, color), (offset+15, 25))
            first_points[0] = [offset+15, 25]
            first_points[3] = [offset+15, 25+90]

            offset += w+15
            for i in range(string_length-1):
                if i == 1:
                    offset +=  22
                offset +=  12
                
                if (i == string_length -2) and text[i+1] in ["领", "警"]:
                    pil_img.paste( GenNewCh(self.fontCE, text[i+1], w, 90, (255, 0 ,0)), (offset, 25))
                else:
                    pil_img.paste( GenNewCh(self.fontCE, text[i+1], w, 90, color), (offset, 25))
                offset += w

                if i == string_length -2:
                    first_points[1] = [offset, 25]
                    first_points[2] = [offset, 25+90]


            pts.append(first_points)
            labels.append(text)

        return pil_img, pts, labels


    def gen_plate_string(self):
        '''
        生成车牌号码字符串
        统一绿牌与蓝牌的字符长度，将绿牌后6位随机去掉一个
        '''
        extra_char = ["挂", "学", "警", "使", "港", "澳", "领", "超", "试", "临"]
        plate_str = ""
        random_copy = np.random.randint(40) == 0
        
        if np.random.randint(10) < 6 :
            char = np.random.choice(PLATE_CHARS_DIGIT) 
        else :
            char = np.random.choice(PLATE_CHARS_LETTER)

        # string_length = 8 if self.plate_type == 1 else 7
        string_length = 7
        # 绿牌改成7位
        random_remove_index = np.random.randint(2,8) if self.plate_type == 1 else -1

        for cpos in range(string_length):
            if cpos == 0:
                tmp_char = copy.copy(PLATE_CHARS_PROVINCE)
                if self.plate_type == 4:
                    tmp_char.extend(["使"] * 20)

                if ENCHANCE_京津冀:
                    indcies_list = list(np.arange(len(tmp_char))) +[0]*5 + [2]*5 + [4]*5 
                    np.random.shuffle(indcies_list)
                    np.random.shuffle(indcies_list)
                    plate_str += tmp_char[indcies_list[0]]
                else:
                    plate_str += np.random.choice(tmp_char)

            elif cpos == 1:
                tmp_char = copy.copy(PLATE_CHARS_LETTER)
                if self.plate_type == 4:
                    tmp_char.extend(PLATE_CHARS_DIGIT)
                plate_str += np.random.choice(tmp_char)
            else:
                if np.random.randint(10) < 6:
                    tmp_char = copy.copy(PLATE_CHARS_DIGIT)
                else :
                    tmp_char = copy.copy(PLATE_CHARS_LETTER)
                # 绿牌 去掉一个数字
                if cpos == random_remove_index:
                    # 使用小写x 代替
                    plate_str += "x" 
                # 挂 学 警
                if cpos == string_length -1 and self.plate_type in [3, 4, 5]:
                    tmp_char.extend(extra_char*3)

                # 黄牌 蓝牌 + ...
                if random_copy and np.random.randint(2) == 0:
                    plate_str += char
                else :
                    plate_str += np.random.choice(tmp_char)
        return plate_str


    def warp_perspective(self, img,max_x_ratio=0.55, max_y_ratio=0.4, max_h_ratio=0.2):
        """ 使图像轻微的畸变
            img 输入图像
            factor 畸变的参数
            size 为图片的目标尺寸
        """
        size_o = [img.shape[1],img.shape[0]]

        pts1 = np.float32([[0,0],[size_o[0],0],[size_o[0],size_o[1]],[0,size_o[1]]])
        
        # 车牌整体水平偏移, update ->  只会变短
        prob = list(np.arange(100)) + list(np.arange(60,100))  + list(np.arange(80,100))
        prob_choice  = np.random.choice(prob)
        # # 车牌右上 右下向左边收缩的偏移
        x_diff = int(float(size_o[0] * prob_choice)* 0.01 *max_x_ratio )
        # # 车牌右上 右下向上边抬的偏移
        y_diff = int(float(size_o[1] * prob_choice )* 0.01  *max_y_ratio )
        # # 车牌左上和右上点 左右摇摆的 偏移
        h_diff = int(float(size_o[1] * prob_choice )* 0.01  *max_h_ratio ) # int(size_o[0]* 0.1) # np.random.randint(int(img.shape[1]* 0.1))
        # 右下点往上抬的偏移
        v_diff = np.random.randint(int(size_o[1]* 0.1), int(size_o[1]* 0.25))

        # 去掉改左边的变换
        if False:
            # change left
            # x_diff = -x_diff if np.random.randint(2) < 1 else x_diff
            # y_diff = -y_diff if np.random.randint(2) < 1 else y_diff
            x_diff = x_diff
            y_diff = -y_diff
            # 水平偏移
            h_diff = np.random.randint(int(img.shape[1]* 0.05))
            if np.random.randint(2) < 1: h_diff = -h_diff

            p1 = [
                    0 ,
                    0 if y_diff < 0 else y_diff]
            p2 = [
                size_o[0]-x_diff+h_diff if x_diff < 0 else size_o[0]-x_diff+h_diff,
                  -y_diff  if y_diff < 0  else 0]
            p3 = [
                size_o[0]-x_diff if x_diff < 0 else size_o[0]-x_diff,
                size_o[1]-y_diff  if y_diff < 0  else size_o[1]
            ]
            # 竖直偏移
            v_diff = np.random.randint(int(img.shape[1]* 0.05))
            p4 =[
                0,
                0 + size_o[1]- v_diff if y_diff < 0 else y_diff + size_o[1] - v_diff
            ]
            pts2 = np.float32([p1,p2,p3,p4])
        else:
            # change right
            x_diff = -x_diff
            y_diff = -y_diff

            if np.random.randint(2) < 1: h_diff = -h_diff

            p1 = [
                    0 if h_diff < 0 else h_diff, 
                    -y_diff if y_diff < 0 else 0]
            p2 = [
                    size_o[0]+x_diff if h_diff < 0 else size_o[0]+x_diff+h_diff,
                    0  if y_diff < 0  else y_diff
                  ]
            p3 = [
                size_o[0]+x_diff-h_diff  if h_diff < 0 else size_o[0]+x_diff,
                  size_o[1] - v_diff if y_diff < 0  else size_o[1] +y_diff- v_diff 
            ]
            p4 =[
                -h_diff if h_diff < 0 else 0, 
                size_o[1] -y_diff if y_diff < 0 else  size_o[1]
            ]
            pts2 = np.float32([p1,p2,p3,p4])

        M  = cv2.getPerspectiveTransform(pts1,pts2)
        max_x ,max_y= np.max(pts2,axis=0 )

        dst = cv2.warpPerspective(img,M,(max_x, max_y))

        return dst, pts2


    def run(self, images_in_per_dir, via_name="via_region_data_ori.json"):
        one_via_file = DATASET_DIR + os.sep + self.current_name + os.sep + via_name
        one_data_dict = filesystem.read_via_file(one_via_file, allow_empty=True)
        image_list = list(one_data_dict.keys())

        epochs = len(image_list) // images_in_per_dir
        for e in range(epochs + 1):
            start = e*images_in_per_dir
            end = (e+1)*images_in_per_dir

            epoch_list = image_list[ start: end ]
            if len(epoch_list) == 0:continue

            new_save_dir = SAVE_DIR + os.sep + self.current_name + "_{}".format(e)
            if os.path.exists(new_save_dir+os.sep+via_name):
                print("continue: ",new_save_dir )
                continue

            save_info = []
            start_img_idx = np.random.randint(len(self.template_image_paths))
            for background_path in tqdm(epoch_list):
                start_img_idx += 1
                # background_path = self.template_image_paths[start_img_idx % len(self.template_image_paths)]
                background_image = cv2.imread(background_path)
                de = data_enhance.DataEnhance(background_image, enlarge_mask=True, 
                    data_dict=one_data_dict[background_path], flag="polygon", 
                    convert_to_rect=False, default_label=None)      # 取消原图上的牌照转换成 plate 因为原图上的拍照不够精确 无法训练DB

                for _ in range(np.random.randint(2,4)):
                    # # 更新背景图
                    if self.plate_type == 1:
                        pil_img = Image.new("RGBA", (428, 140), (0, 0, 0,0))
                    elif self.plate_type in [0, 2, 4, 5]:
                        pil_img = Image.new("RGBA", (440, 140), (0, 0, 0,0))
                    elif self.plate_type == 3:
                        pil_img = Image.new("RGBA", (440, 220), (0, 0, 0,0))
                    else:
                        raise NotImplementedError
                    plate_str = self.gen_plate_string()
                    text_line_image, pts, labels = self.gen_text_line_image(pil_img, plate_str) 
                    # print(pts)
                    # return 0

                    background = Image.open(self.template_image_paths[start_img_idx % len(self.template_image_paths)]) # 060086
                    background.paste(text_line_image, (0,0), text_line_image)

                    ratio = np.random.choice(list(list(range(42,50)) + list(range(42,60)) + list(range(42,70)) )) * 0.01
                    text_image = background.resize((int(background.size[0] * ratio), int(background.size[1] *  ratio)), Image.ANTIALIAS)
                    pts = np.array(pts, dtype=np.int) * ratio


                    cv_image, mask, new_pts3, new_labels = imgaug_tool.aug_image_for_plate_v2(text_image, pts, labels)
                    if type(cv_image) == type(None):
                        continue 

                    cv_image = imgaug_tool.aug_image_for_plate(cv_image)

                    # aug_img_o ,pts2= self.warp_perspective(aug_img_o)
                    # array_image, pts3, image_mask = imgaug_tool.random_add_angle(aug_img_o,pts2, 3)

                    ok = de.random_add_region(cv_image, new_labels, polygon=new_pts3, mask=mask)
                save_info.append(de)

            # save
            os.makedirs(new_save_dir, exist_ok=True)
            save_dict  =dict()
            for idx, de in enumerate(save_info):
                file_name = "{}_{}.jpg".format(e, idx)
                save_dict[file_name] = de.convert_to_dict(new_save_dir + os.sep+ file_name)

            with open(new_save_dir + os.sep + via_name, "w") as wf:
                wf.write(json.dumps(save_dict))

# 生成蓝牌背景图
def gen_plate_background():
    template_image = "/media/swls/disk1/vanlance/project/plate_recognition/font/template.bmp"
    # 069BEA   4B677F  075B99 10A3E8 153661 3393C3
    # 052bdf
    save_path = "/media/swls/disk1/vanlance/project/plate_recognition/font/template"
    r_min = 16*0+ 6
    r_max = 16*4+11

    g_min = 16*3+6
    g_max = 16*10+3

    b_min = 16*6+1
    b_max = 16*14+8

    img = cv2.imread(template_image) # 060086
    for r in range(r_min, r_max, 3):
        for g in range(g_min, g_max, 6):
            for b in range(b_min, b_max, 6):
                array_img = copy.copy(img)
                array_img[array_img == 6] = r
                array_img[array_img == 0] = g
                array_img[array_img == (8*16+6)] = b

                cv2.imwrite(save_path +os.sep + "{}_{}_{}.bmp".format(r,g,b), array_img)

    # array_img[array_img == 6] = 26
    # array_img[array_img == 0] = 122
    # array_img[array_img == (8*16+6)] = 219
    # cv2.imwrite("_2.".join(TEMPLATE_IMAGE.split(".")), array_img)

# 生成绿牌背景图
def gen_new_green_plate_background():
    template_image = "/media/swls/disk1/vanlance/project/plate_recognition/font/3.png"
    # 069BEA   4B677F  075B99 10A3E8 153661 3393C3
    # 052bdf
    save_path = "/media/swls/disk1/vanlance/project/plate_recognition/font/new_green_template"
    r_min = -30
    r_max = 30

    g_min = -60
    g_max = 0

    b_min = -30
    b_max = 30

    img = cv2.imread(template_image) # 060086
    for r in range(r_min, r_max, 5):
        for g in range(g_min, g_max, 6):
            for b in range(b_min, b_max, 6):
                array_img = copy.copy(img)
                array_img[:,:,2] += np.array(r, np.uint8)
                array_img[:,:,1] += np.array(g, np.uint8)
                array_img[:,:,0] += np.array(b, np.uint8)

                cv2.imwrite(save_path +os.sep + "3_{}_{}_{}.png".format(r,g,b), array_img)

    # array_img[array_img == 6] = 26
    # array_img[array_img == 0] = 122
    # array_img[array_img == (8*16+6)] = 219
    # cv2.imwrite("_2.".join(TEMPLATE_IMAGE.split(".")), array_img)

# 生成黄牌背景图
def gen_yellow_plate_background():
    template_image = "/media/swls/disk1/vanlance/project/plate_recognition/font/20120625161654043_c.jpg"
    # 069BEA   4B677F  075B99 10A3E8 153661 3393C3
    # 052bdf
    save_path = "/media/swls/disk1/vanlance/project/plate_recognition/font/yellow_template"
    
    # 255
    r_min = -85
    r_max = 0
    
    # 190
    g_min = -70
    g_max = 0

    # 0
    b_min = 0
    b_max = 57

    img = cv2.imread(template_image) # 060086
    for r in range(r_min, r_max, 5):
        for g in range(g_min, g_max, 6):
            for b in range(b_min, b_max, 6):
                array_img = copy.copy(img)
                array_img[:,:,2] += np.array(r, np.uint8)
                array_img[:,:,1] += np.array(g, np.uint8)
                array_img[:,:,0] += np.array(b, np.uint8)

                cv2.imwrite(save_path +os.sep + "10_3_{}_{}_{}.png".format(r,g,b), array_img)

    # array_img[array_img == 6] = 26
    # array_img[array_img == 0] = 122
    # array_img[array_img == (8*16+6)] = 219
    # cv2.imwrite("_2.".join(TEMPLATE_IMAGE.split(".")), array_img)

# 生成双层黄牌背景图
def gen_double_yellow_plate_background():
    template_image = "/mnt/disk1/vanlance/project/plate_recognition/font/double_plate_white.png"
    # 069BEA   4B677F  075B99 10A3E8 153661 3393C3
    # 052bdf
    save_path = "/mnt/disk1/vanlance/project/plate_recognition/font/template_yellow_double"
    
    # 255
    r_min = -85
    r_max = 0
    
    # 190
    g_min = -70
    g_max = 0

    # 0
    b_min = 0
    b_max = 57

    img = cv2.imread(template_image) # 060086
    for r in range(r_min, r_max, 5):
        for g in range(g_min, g_max, 6):
            for b in range(b_min, b_max, 6):
                array_img = copy.copy(img)
                array_img[:,:,2] += np.array(r, np.uint8)
                array_img[:,:,1] += np.array(g, np.uint8)
                array_img[:,:,0] += np.array(b, np.uint8)

                cv2.imwrite(save_path +os.sep + "11_3_{}_{}_{}.png".format(r,g,b), array_img)

    # array_img[array_img == 6] = 26
    # array_img[array_img == 0] = 122
    # array_img[array_img == (8*16+6)] = 219
    # cv2.imwrite("_2.".join(TEMPLATE_IMAGE.split(".")), array_img)

# 生成黑牌背景图
def gen_black_plate_background():
    template_image = "/mnt/disk1/vanlance/project/plate_recognition/font/plate_black0.png"
    # 069BEA   4B677F  075B99 10A3E8 153661 3393C3
    # 052bdf
    save_path = "/mnt/disk1/vanlance/project/plate_recognition/font/template_black"
    
    # 255
    r_min = 0
    r_max = 40
    
    # 190
    g_min = 0
    g_max = 40

    # 0
    b_min = 0
    b_max = 40

    img = cv2.imread(template_image) # 060086
    for r in range(r_min, r_max, 3):
        for g in range(g_min, g_max, 3):
            for b in range(b_min, b_max, 3):
                array_img = copy.copy(img)
                array_img[:,:,2] += np.array(r, np.uint8)
                array_img[:,:,1] += np.array(g, np.uint8)
                array_img[:,:,0] += np.array(b, np.uint8)

                cv2.imwrite(save_path +os.sep + "11_3_{}_{}_{}.png".format(r,g,b), array_img)

# 生成白牌背景图
def gen_white_plate_background():
    template_image = "/mnt/disk1/vanlance/project/plate_recognition/font/plate_white_1.png"
    # 069BEA   4B677F  075B99 10A3E8 153661 3393C3
    # 052bdf
    save_path = "/mnt/disk1/vanlance/project/plate_recognition/font/template_white"
    
    # 255
    r_min = -40
    r_max = 0
    
    # 190
    g_min = -40
    g_max = 0

    # 0
    b_min = -40
    b_max = 0

    img = cv2.imread(template_image) # 060086
    for r in range(r_min, r_max, 3):
        for g in range(g_min, g_max, 3):
            for b in range(b_min, b_max, 3):
                array_img = copy.copy(img)
                array_img[:,:,2] += np.array(r, np.uint8)
                array_img[:,:,1] += np.array(g, np.uint8)
                array_img[:,:,0] += np.array(b, np.uint8)

                cv2.imwrite(save_path +os.sep + "11_3_{}_{}_{}.png".format(r,g,b), array_img)

def add_four_anchor_for_plate(image_dir, text_color="#B0B0B0,#F0F0F0", text_color2="#000000,#404040"):
    images = [image_dir + os.sep + p for p in os.listdir(image_dir)]
    np.random.shuffle(images)
    np.random.shuffle(images)

    for img_p in images[:len(images) // 2] :
        cv_img = cv2.imread(img_p)
        # colors = [ImageColor.getrgb(c) for c in text_color2.split(',')]
        # c1, c2 = colors[0], colors[-1]
        # fill = (
        #     random.randint(c1[0], c2[0]),
        #     random.randint(c1[1], c2[1]),
        #     random.randint(c1[2], c2[2])
        # )
        radis = np.random.randint(10,15)

        colors1 = [ImageColor.getrgb(c) for c in text_color2.split(',')]
        c1_1, c2_1 = colors1[0], colors1[-1]
        fill_1 = ( random.randint(c1_1[0], c2_1[0]), random.randint(c1_1[1], c2_1[1]), random.randint(c1_1[2], c2_1[2]) )
        colors2 = [ImageColor.getrgb(c) for c in text_color.split(',')]
        c1_2, c2_2 = colors2[0], colors2[-1]
        fill_2 = ( random.randint(c1_2[0], c2_2[0]), random.randint(c1_2[1], c2_2[1]), random.randint(c1_2[2], c2_2[2]) )

        if np.random.randint(2):
            fill = fill_1
        else :
            fill = fill_2

        # # double yellow
        # cv2.circle(cv_img, (170, 25), radis // 2 +1, fill, radis)
        # cv2.circle(cv_img, (200, 25), radis // 2 +1, fill, radis)
        # cv2.circle(cv_img, (680, 415), radis // 2 +1, fill, radis)
        # cv2.circle(cv_img, (710, 415), radis // 2 +1, fill, radis)

        # if np.random.randint(2):
        #     fill = fill_1
        # else :
        #     fill = fill_2
        # cv2.circle(cv_img, (680, 25), radis // 2 +1, fill, radis)
        # cv2.circle(cv_img, (710, 25), radis // 2 +1, fill, radis)
        # cv2.circle(cv_img, (170, 415), radis // 2 +1, fill, radis)
        # cv2.circle(cv_img, (200, 415), radis // 2 +1, fill, radis)

        # colors = [ImageColor.getrgb(c) for c in text_color2.split(',')]
        # c1, c2 = colors[0], colors[-1]
        # fill3 = ( random.randint(c1[0], c2[0]), random.randint(c1[1], c2[1]), random.randint(c1[2], c2[2]) )
        # cv2.circle(cv_img, (440, 90), radis // 2 +1, fill3, radis)

        # # yellow
        cv2.circle(cv_img, (108, 12), radis // 2 +1, fill, radis)
        cv2.circle(cv_img, (317, 12), radis // 2 +1, fill, radis)
        cv2.circle(cv_img, (108, 127), radis // 2 +1, fill, radis)
        cv2.circle(cv_img, (317, 127), radis // 2 +1, fill, radis)
        
        # # green plate
        # cv2.circle(cv_img, (124, 12), radis // 2 +1, fill, radis)
        # cv2.circle(cv_img, (304, 12), radis // 2 +1, fill, radis)
        # cv2.circle(cv_img, (124, 127), radis // 2 +1, fill, radis)
        # cv2.circle(cv_img, (304, 127), radis // 2 +1, fill, radis)        
        cv2.imwrite(img_p, cv_img)
        # return

def crop_image(data_dir, save_dir):
    for p in os.listdir(data_dir):
        image_path = data_dir  + os.sep + p 
        image = cv2.imread(image_path)
        save_image = np.concatenate((image[:, 0:220], image[:, 272:480]), axis=1)
        save_path = save_dir + os.sep + os.path.basename(image_path)
        cv2.imwrite(save_path, save_image)

def convert_to_rect_points(data_dir, via_name="via_region_data.json"):
    with open(data_dir + os.sep + via_name) as rf:
        data_dict = json.loads(rf.read())

    new_data_dict = dict()
    for k in data_dict.keys():
        item = data_dict[k]
        regions = via_tool.read_via_polygon_to_rect_imp(item["regions"])
        for pts, label in regions:
            if label == "retain":continue
            region = dict()
            shape_attributes = dict()
            shape_attributes["name"] = "polygon"
            all_points_x = []
            all_points_y = []
            for xy in pts:
                all_points_x.append(int(xy[0]))
                all_points_y.append(int(xy[1]))
            shape_attributes["all_points_x"] = all_points_x
            shape_attributes["all_points_y"] = all_points_y
            region["shape_attributes"] = shape_attributes
            region["region_attributes"] = {"label": "plate"}
            item["regions"].append(region)
        new_data_dict[k] = item

    with open(data_dir + os.sep + "via_region_data_ori.json", "w") as wf:
        wf.write(json.dumps(new_data_dict))

def thread_run(images_in_per_dir, current_name, current_type):
    tig = TextImageGenerator(current_name, current_type)
    tig.run(images_in_per_dir)


if __name__ == "__main__":

    total_count = 360           # via文件夹的个数
    # 0 -> 蓝牌  1 -> 绿牌  2 -> 黄牌 3 -> 双层黄牌 4 -> 黑牌  5 -> 白牌
    total_type = [0, 1, 2, 3, 4, 5]   # 车牌类型总数
    images_in_per_dir = 500

    thread_count = 4
    ENCHANCE_京津冀 = False

    DATASET_DIR = "/mnt/disk2/datasets_kanggle/ocr/plate_data_CCPD/plate_via"
    SAVE_DIR = "/home/swls/work_dir/github/paddle/train_data/det/plate/images"

    TEMPLATE_IMAGES = [
        "/mnt/disk1/vanlance/project/plate_recognition/font/template_blue",
        "/mnt/disk1/vanlance/project/plate_recognition/font/template_green2",
        "/mnt/disk1/vanlance/project/plate_recognition/font/template_yellow",
        "/mnt/disk1/vanlance/project/plate_recognition/font/template_yellow_double2",
        "/mnt/disk1/vanlance/project/plate_recognition/font/template_black",
        "/mnt/disk1/vanlance/project/plate_recognition/font/template_white"
    ]
    fontChEng = [
        "/mnt/disk1/vanlance/project/plate_recognition/font/plate_blue.otf",
        "/mnt/disk1/vanlance/project/plate_recognition/font/vanlance.ttf",
        "/mnt/disk1/vanlance/project/plate_recognition/font/plate_blue.otf",
        "/mnt/disk1/vanlance/project/plate_recognition/font/plate_blue.otf",
        "/mnt/disk1/vanlance/project/plate_recognition/font/plate_blue.otf",
        "/mnt/disk1/vanlance/project/plate_recognition/font/plate_blue.otf"
    ]

    my_type = [0] * 80 + \
              [1] * 80 + \
              [2] * 80 + \
              [3] * 80 + \
              [4] * 20 + \
              [5] * 20 
    np.random.shuffle(my_type)
    np.random.shuffle(my_type)

    # for idx, name in enumerate(os.listdir(DATASET_DIR)):
    #     thread_run(images_in_per_dir, name, my_type[idx])   # np.random.choice(total_type)


    # thread_count = 3
    # p = Pool(thread_count)
    # for idx, name in enumerate(os.listdir(DATASET_DIR)):
    #     p.apply_async(thread_run, args=(images_in_per_dir, name, my_type[idx] ))
    # p.close()
    # p.join()
    # print('All subprocesses done.')


    # # 1. 生成车牌
    # gen_plate_background()
    # gen_new_green_plate_background()   # new_green_template
    # gen_yellow_plate_background()
    # gen_double_yellow_plate_background()
    # gen_black_plate_background()
    # gen_white_plate_background

    # # 2. 制作车牌背景图片 绿牌
    # data_dir = "/media/swls/disk1/vanlance/project/plate_recognition/font/new_green_template"
    # save_dir = "/media/swls/disk1/vanlance/project/plate_recognition/font/new_green_template_2"
    # crop_image(data_dir, save_dir)

    # # 在车牌背景图片上增加四个瞄点
    # TEMPLATE_IMAGES = "/mnt/disk1/vanlance/project/plate_recognition/font/template_black"
    # TEMPLATE_IMAGES = "/media/swls/disk1/vanlance/project/plate_recognition/font/yellow_template"
    # # TEMPLATE_IMAGES = "/media/swls/disk1/vanlance/project/plate_recognition/font/new_green_template_2"
    # add_four_anchor_for_plate(TEMPLATE_IMAGES)

    ## 将标记的轮廓转化为矩形 四点
    # data_dir = "/home/swls/work_dir/github/paddle/train_data/det/plate/images/0_0_0_3_via"
    # convert_to_rect_points(data_dir, "via_region_data.json")

"""

cd /home/swls/work_dir/git/python_script && conda activate tf3 && export PYTHONPATH=$PYTHONPATH:`pwd` && sh sunjie/plate_recognition/gen_plate_0.sh

@2020-01-08     修改 warp_perspective 函数默认参数
                修改车牌背景图片颜色选取方式
                增强车牌相似字符
"""